#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Kernel Capabilities Management
内核能力管理模块

该模块负责：
1. 内核特性的运行时检测
2. 能力配置的加载和管理
3. 策略选择的智能化决策
"""

from __future__ import print_function
from __future__ import unicode_literals
from __future__ import division
from __future__ import absolute_import

import os
import sys
import json
import ctypes
import subprocess
from collections import defaultdict

# Python 2/3 兼容性
if sys.version_info[0] == 2:
    string_types = basestring
    text_type = unicode
else:
    string_types = str
    text_type = str


class KernelCapabilities(object):
    """内核能力管理器"""
    
    def __init__(self, config_path=None):
        self.config_path = config_path or self._get_default_config_path()
        self.capabilities = {}
        self.runtime_checks = {}
        self.load_capabilities()
    
    def _get_default_config_path(self):
        """获取默认配置文件路径"""
        return os.path.join(
            os.path.dirname(__file__), '..', '..', 
            'config', 'kernel_capabilities.json'
        )
    
    def load_capabilities(self):
        """加载能力配置"""
        if os.path.exists(self.config_path):
            try:
                with open(self.config_path, 'r') as f:
                    data = json.load(f)
                    self.capabilities = data.get('capabilities', {})
                    print(f"[+] 已加载内核能力配置: {self.config_path}")
            except Exception as e:
                print(f"[!] 加载配置失败: {e}")
                self._initialize_default_capabilities()
        else:
            print("[*] 配置文件不存在，使用默认配置")
            self._initialize_default_capabilities()
    
    def _initialize_default_capabilities(self):
        """初始化默认能力配置"""
        self.capabilities = {
            'memfd_create': False,
            'memfd_secret': False,
            'bpf_support': False,
            'ebpf_support': False,
            'seccomp_support': False,
            'namespace_support': False,
            'cgroup_v2': False,
            'user_namespaces': False,
            'mount_namespaces': False,
            'pid_namespaces': False,
            'net_namespaces': False,
            'uts_namespaces': False,
            'ipc_namespaces': False
        }
    
    def check_memfd_capabilities(self):
        """检查内存文件描述符能力"""
        results = {
            'memfd_create': False,
            'memfd_secret': False,
            'memfd_create_available': False,
            'memfd_secret_available': False
        }
        
        try:
            # 检查 memfd_create
            libc = ctypes.CDLL("libc.so.6")
            
            # 尝试获取 memfd_create 函数
            if hasattr(libc, 'memfd_create'):
                results['memfd_create'] = True
                results['memfd_create_available'] = True
                
                # 尝试实际调用 (测试用)
                try:
                    # memfd_create(name, flags)
                    # MFD_CLOEXEC = 1
                    fd = libc.memfd_create(b"test", 1)
                    if fd > 0:
                        os.close(fd)
                        results['memfd_create_available'] = True
                except:
                    results['memfd_create_available'] = False
            
            # 检查 memfd_secret (Linux 5.14+)
            if hasattr(libc, 'memfd_secret'):
                results['memfd_secret'] = True
                try:
                    # 尝试调用 memfd_secret
                    fd = libc.memfd_secret(0)
                    if fd > 0:
                        os.close(fd)
                        results['memfd_secret_available'] = True
                except:
                    results['memfd_secret_available'] = False
                    
        except Exception as e:
            print(f"[!] memfd 能力检查失败: {e}")
        
        self.runtime_checks.update(results)
        return results
    
    def check_bpf_capabilities(self):
        """检查 BPF 能力"""
        results = {
            'bpf_support': False,
            'ebpf_support': False,
            'bpf_runtime_available': False,
            'bpf_jit_enabled': False,
            'bpf_unprivileged_enabled': False
        }
        
        try:
            # 检查 BPF 文件系统
            if os.path.exists('/sys/fs/bpf'):
                results['bpf_support'] = True
                results['bpf_runtime_available'] = True
            
            # 检查 BPF 统计信息
            if os.path.exists('/proc/sys/kernel/bpf_stats_enabled'):
                results['ebpf_support'] = True
            
            # 检查 BPF JIT 编译
            jit_path = '/proc/sys/net/core/bpf_jit_enable'
            if os.path.exists(jit_path):
                try:
                    with open(jit_path, 'r') as f:
                        jit_value = f.read().strip()
                        results['bpf_jit_enabled'] = jit_value == '1'
                except:
                    pass
            
            # 检查非特权 BPF
            unpriv_path = '/proc/sys/kernel/unprivileged_bpf_disabled'
            if os.path.exists(unpriv_path):
                try:
                    with open(unpriv_path, 'r') as f:
                        unpriv_value = f.read().strip()
                        results['bpf_unprivileged_enabled'] = unpriv_value == '0'
                except:
                    pass
            
            # 尝试加载简单的 BPF 程序进行测试
            if results['bpf_support']:
                results.update(self._test_bpf_program_loading())
                
        except Exception as e:
            print(f"[!] BPF 能力检查失败: {e}")
        
        self.runtime_checks.update(results)
        return results
    
    def _test_bpf_program_loading(self):
        """测试 BPF 程序加载能力"""
        test_results = {
            'bpf_program_loading': False,
            'bpf_map_creation': False
        }
        
        try:
            # 尝试使用 bcc 库进行测试
            try:
                from bcc import BPF
                
                # 简单的 BPF 程序
                prog = """
                int hello(void *ctx) {
                    return 0;
                }
                """
                
                b = BPF(text=prog)
                test_results['bpf_program_loading'] = True
                test_results['bpf_map_creation'] = True
                
            except ImportError:
                # 如果没有 bcc，尝试直接系统调用
                test_results.update(self._test_bpf_syscall())
                
        except Exception as e:
            print(f"[!] BPF 程序测试失败: {e}")
        
        return test_results
    
    def _test_bpf_syscall(self):
        """测试 BPF 系统调用"""
        results = {
            'bpf_program_loading': False,
            'bpf_map_creation': False
        }
        
        try:
            # 使用 ctypes 调用 bpf 系统调用
            libc = ctypes.CDLL("libc.so.6")
            
            # BPF 系统调用号 (x86_64: 321)
            SYS_bpf = 321
            
            # 尝试创建 BPF map
            # 这里只是测试系统调用是否可用，不实际创建
            # 实际实现需要更复杂的参数构造
            
        except Exception as e:
            print(f"[!] BPF 系统调用测试失败: {e}")
        
        return results
    
    def check_namespace_capabilities(self):
        """检查命名空间能力"""
        results = {
            'user_namespaces': False,
            'mount_namespaces': False,
            'pid_namespaces': False,
            'net_namespaces': False,
            'uts_namespaces': False,
            'ipc_namespaces': False,
            'cgroup_namespaces': False
        }
        
        try:
            # 检查 /proc/self/ns/ 目录
            ns_dir = '/proc/self/ns'
            if os.path.exists(ns_dir):
                ns_files = os.listdir(ns_dir)
                
                namespace_map = {
                    'user': 'user_namespaces',
                    'mnt': 'mount_namespaces',
                    'pid': 'pid_namespaces',
                    'net': 'net_namespaces',
                    'uts': 'uts_namespaces',
                    'ipc': 'ipc_namespaces',
                    'cgroup': 'cgroup_namespaces'
                }
                
                for ns_file, capability in namespace_map.items():
                    if ns_file in ns_files:
                        results[capability] = True
            
            # 检查用户命名空间的特殊限制
            if results['user_namespaces']:
                results.update(self._check_user_namespace_restrictions())
                
        except Exception as e:
            print(f"[!] 命名空间能力检查失败: {e}")
        
        self.runtime_checks.update(results)
        return results
    
    def _check_user_namespace_restrictions(self):
        """检查用户命名空间限制"""
        restrictions = {
            'user_ns_unprivileged_enabled': True,
            'user_ns_max_count': 0
        }
        
        try:
            # 检查是否允许非特权用户创建用户命名空间
            unpriv_path = '/proc/sys/kernel/unprivileged_userns_clone'
            if os.path.exists(unpriv_path):
                with open(unpriv_path, 'r') as f:
                    value = f.read().strip()
                    restrictions['user_ns_unprivileged_enabled'] = value == '1'
            
            # 检查最大用户命名空间数量
            max_path = '/proc/sys/user/max_user_namespaces'
            if os.path.exists(max_path):
                with open(max_path, 'r') as f:
                    value = f.read().strip()
                    restrictions['user_ns_max_count'] = int(value)
                    
        except Exception as e:
            print(f"[!] 用户命名空间限制检查失败: {e}")
        
        return restrictions
    
    def check_seccomp_capabilities(self):
        """检查 seccomp 能力"""
        results = {
            'seccomp_support': False,
            'seccomp_filter': False,
            'seccomp_user_notif': False
        }
        
        try:
            # 检查 seccomp 支持
            if os.path.exists('/proc/sys/kernel/seccomp'):
                results['seccomp_support'] = True
                
                # 检查过滤器支持
                try:
                    import ctypes
                    libc = ctypes.CDLL("libc.so.6")
                    
                    # PR_SET_SECCOMP = 22
                    # SECCOMP_MODE_FILTER = 2
                    if hasattr(libc, 'prctl'):
                        results['seccomp_filter'] = True
                        
                except:
                    pass
            
            # 检查用户通知支持 (Linux 5.0+)
            seccomp_dir = '/proc/sys/kernel'
            if os.path.exists(os.path.join(seccomp_dir, 'seccomp')):
                # 这是一个简化的检查，实际需要更复杂的测试
                results['seccomp_user_notif'] = True
                
        except Exception as e:
            print(f"[!] seccomp 能力检查失败: {e}")
        
        self.runtime_checks.update(results)
        return results
    
    def check_cgroup_capabilities(self):
        """检查 cgroup 能力"""
        results = {
            'cgroup_v1': False,
            'cgroup_v2': False,
            'cgroup_unified': False
        }
        
        try:
            # 检查 cgroup v1
            if os.path.exists('/sys/fs/cgroup'):
                cgroup_contents = os.listdir('/sys/fs/cgroup')
                if any(item for item in cgroup_contents if not item.startswith('.')):
                    results['cgroup_v1'] = True
            
            # 检查 cgroup v2
            cgroup_v2_path = '/sys/fs/cgroup/cgroup.controllers'
            if os.path.exists(cgroup_v2_path):
                results['cgroup_v2'] = True
                
                # 检查是否为统一层次结构
                try:
                    with open(cgroup_v2_path, 'r') as f:
                        controllers = f.read().strip()
                        if controllers:
                            results['cgroup_unified'] = True
                except:
                    pass
                    
        except Exception as e:
            print(f"[!] cgroup 能力检查失败: {e}")
        
        self.runtime_checks.update(results)
        return results
    
    def run_full_capability_check(self):
        """运行完整的能力检查"""
        print("[*] 开始运行时能力检查...")
        
        checks = [
            ('内存文件描述符', self.check_memfd_capabilities),
            ('BPF支持', self.check_bpf_capabilities),
            ('命名空间', self.check_namespace_capabilities),
            ('seccomp', self.check_seccomp_capabilities),
            ('cgroup', self.check_cgroup_capabilities)
        ]
        
        for check_name, check_func in checks:
            try:
                print(f"[*] 检查 {check_name}...")
                check_func()
            except Exception as e:
                print(f"[!] {check_name} 检查失败: {e}")
        
        print(f"[+] 运行时检查完成，共检测 {len(self.runtime_checks)} 项能力")
        return self.runtime_checks
    
    def get_capability_status(self, capability_name):
        """获取特定能力状态"""
        # 优先返回运行时检查结果
        if capability_name in self.runtime_checks:
            return self.runtime_checks[capability_name]
        
        # 回退到配置文件结果
        return self.capabilities.get(capability_name, False)
    
    def is_feature_available(self, feature_name):
        """检查特性是否可用"""
        return self.get_capability_status(feature_name)
    
    def get_available_features(self):
        """获取所有可用特性列表"""
        available = []
        
        # 合并配置文件和运行时检查结果
        all_capabilities = {}
        all_capabilities.update(self.capabilities)
        all_capabilities.update(self.runtime_checks)
        
        for feature, available_status in all_capabilities.items():
            if available_status:
                available.append(feature)
        
        return available
    
    def generate_capability_report(self):
        """生成能力报告"""
        report = {
            'summary': {
                'total_capabilities': len(self.capabilities),
                'runtime_checks': len(self.runtime_checks),
                'available_features': len(self.get_available_features())
            },
            'capabilities': self.capabilities,
            'runtime_checks': self.runtime_checks,
            'available_features': self.get_available_features()
        }
        
        return report
    
    def print_capability_summary(self):
        """打印能力摘要"""
        available_features = self.get_available_features()
        
        print("\n" + "="*50)
        print("内核能力检查报告")
        print("="*50)
        print(f"配置文件能力: {len(self.capabilities)}")
        print(f"运行时检查: {len(self.runtime_checks)}")
        print(f"可用特性: {len(available_features)}")
        
        print("\n可用特性列表:")
        for feature in sorted(available_features):
            print(f"  ✓ {feature}")
        
        # 显示不可用的特性
        all_features = set(self.capabilities.keys()) | set(self.runtime_checks.keys())
        unavailable = all_features - set(available_features)
        
        if unavailable:
            print("\n不可用特性:")
            for feature in sorted(unavailable):
                print(f"  ✗ {feature}")
        
        print("="*50)


def main():
    """主函数 - 用于测试"""
    capabilities = KernelCapabilities()
    capabilities.run_full_capability_check()
    capabilities.print_capability_summary()


if __name__ == "__main__":
    main()